// Package main 提供了一个异步PDF生成服务
package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/chromedp/cdproto/page"
	"github.com/chromedp/chromedp"
)

// TaskStatus 定义任务状态的类型
type TaskStatus string

// 定义任务的各种状态常量
const (
	TaskStatusPending   TaskStatus = "pending"   // 等待处理
	TaskStatusRunning   TaskStatus = "running"   // 正在处理
	TaskStatusCompleted TaskStatus = "completed" // 处理完成
	TaskStatusFailed    TaskStatus = "failed"    // 处理失败
)

// Task 定义PDF生成任务的结构
type Task struct {
	ID        string     `json:"id"`              // 任务唯一标识符
	Request   PDFRequest `json:"request"`         // PDF生成请求的详细信息
	Status    TaskStatus `json:"status"`          // 任务当前状态
	Error     string     `json:"error,omitempty"` // 错误信息（如果有）
	CreatedAt time.Time  `json:"createdAt"`       // 任务创建时间
	UpdatedAt time.Time  `json:"updatedAt"`       // 任务最后更新时间
}

// TaskManager 管理PDF生成任务的结构
type TaskManager struct {
	tasks    map[string]*Task // 存储所有任务的映射
	tasksMux sync.RWMutex     // 用于任务操作的读写锁
}

// NewTaskManager 创建并初始化一个新的任务管理器
func NewTaskManager() *TaskManager {
	return &TaskManager{
		tasks: make(map[string]*Task),
	}
}

// 全局任务管理器实例
var taskManager *TaskManager

// PDFRequest 定义PDF生成请求的结构
type PDFRequest struct {
	URL      string           `json:"url"`      // 要生成PDF的网页URL
	Filename string           `json:"filename"` // 生成的PDF文件名
	Options  *PDFGenerateOpts `json:"options"`  // PDF生成的选项
}

// PDFGenerateOpts 定义PDF生成的选项
type PDFGenerateOpts struct {
	Landscape           bool    `json:"landscape"`           // 是否横向打印
	DisplayHeaderFooter bool    `json:"displayHeaderFooter"` // 是否显示页眉页脚
	PrintBackground     bool    `json:"printBackground"`     // 是否打印背景
	Scale               float64 `json:"scale"`               // 缩放比例
	PaperWidth          float64 `json:"paperWidth"`          // 纸张宽度（英寸）
	PaperHeight         float64 `json:"paperHeight"`         // 纸张高度（英寸）
	MarginTop           float64 `json:"marginTop"`           // 上边距（英寸）
	MarginBottom        float64 `json:"marginBottom"`        // 下边距（英寸）
	MarginLeft          float64 `json:"marginLeft"`          // 左边距（英寸）
	MarginRight         float64 `json:"marginRight"`         // 右边距（英寸）
}

// PDFResponse 定义API响应的结构
type PDFResponse struct {
	Status   string `json:"status"`             // 响应状态
	Message  string `json:"message,omitempty"`  // 响应消息
	Filename string `json:"filename,omitempty"` // 生成的文件名
	Error    string `json:"error,omitempty"`    // 错误信息
}

// 打印样式常量，用于确保PDF生成时的样式正确
const printStyles = `
@media print {
    /* 确保所有颜色调整准确 */
    * {
        -webkit-print-color-adjust: exact !important;
        print-color-adjust: exact !important;
        color-adjust: exact !important;
        font-family: "WenQuanYi Micro Hei", "Microsoft YaHei", "SimSun", sans-serif !important;
    }
    
    /* 基本页面设置 */
    body {
        width: 100% !important;
        margin: 0 !important;
        padding: 0 !important;
        line-height: 1.45;
        font-family: "WenQuanYi Micro Hei", "Microsoft YaHei", "SimSun", sans-serif !important;
    }

    /* 防止元素在打印时被分页 */
    div, section, header, footer {
        break-inside: avoid;
        page-break-inside: avoid;
    }

    /* 确保背景色正确打印 */
    [style*="background"] {
        -webkit-print-color-adjust: exact !important;
        print-color-adjust: exact !important;
    }

    /* 优化文本渲染 */
    * {
        text-rendering: optimizeLegibility;
        -webkit-font-smoothing: antialiased;
        -moz-osx-font-smoothing: grayscale;
    }

    /* 确保图片和图标正确打印 */
    svg, img, i[class*="fa"] {
        print-color-adjust: exact !important;
        -webkit-print-color-adjust: exact !important;
    }

    /* 确保字体图标正确显示 */
    .fas, .fab, .far {
        font-family: "Font Awesome 5 Free" !important;
        -webkit-print-color-adjust: exact !important;
        color-adjust: exact !important;
    }

    /* 确保渐变色正确显示 */
    [class*="gradient"], [style*="gradient"] {
        -webkit-print-color-adjust: exact !important;
        print-color-adjust: exact !important;
    }
}`

// 默认的PDF生成选项
var defaultPDFOpts = &PDFGenerateOpts{
	Landscape:           false,
	DisplayHeaderFooter: false,
	PrintBackground:     true,
	Scale:               1.0,
	PaperWidth:          8.27, // A4宽度（英寸）
	PaperHeight:         11.7, // A4高度（英寸）
	MarginTop:           0.4,  // 10mm转换为英寸
	MarginBottom:        0.4,
	MarginLeft:          0.4,
	MarginRight:         0.4,
}

// main 函数是程序的入口点
func main() {
	// 初始化任务管理器
	taskManager = NewTaskManager()

	// 创建输出目录
	if err := os.MkdirAll("output", 0755); err != nil {
		log.Fatal("Failed to create output directory:", err)
	}

	// 设置HTTP路由
	http.HandleFunc("/generate-pdf", handlePDFGeneration)
	http.HandleFunc("/download-pdf/", handlePDFDownload)
	http.HandleFunc("/task-status/", handleTaskStatus)

	// 启动HTTP服务器
	port := ":8000"
	fmt.Printf("Server starting on port %s...\n", port)
	log.Fatal(http.ListenAndServe(port, nil))
}

// generatePDF 将指定URL的网页生成PDF文件
func generatePDF(url, outputPath string, opts *PDFGenerateOpts) error {
	// 设置Chrome启动选项
	options := []chromedp.ExecAllocatorOption{
		chromedp.NoFirstRun,
		chromedp.NoDefaultBrowserCheck,
		chromedp.Headless,
		chromedp.DisableGPU,
		chromedp.NoSandbox,
		chromedp.Flag("disable-dev-shm-usage", true),
		chromedp.Flag("font-render-hinting", "none"),
		chromedp.Flag("enable-font-antialiasing", true),
		chromedp.Flag("force-color-profile", "srgb"),
		// 解决cookie partition错误
		chromedp.Flag("disable-features", "CookiePartitioning,PartitionedCookies"),
	}

	// 创建Chrome实例
	ctx, cancel := chromedp.NewExecAllocator(context.Background(), options...)
	defer cancel()

	// 创建新的Chrome上下文
	ctx, cancel = chromedp.NewContext(ctx, chromedp.WithLogf(log.Printf))
	defer cancel()

	// 设置超时时间
	ctx, cancel = context.WithTimeout(ctx, 30*time.Second)
	defer cancel()

	// 生成PDF
	var pdfBuffer []byte
	if err := chromedp.Run(ctx,
		// 导航到目标URL
		chromedp.Navigate(url),
		// 等待页面加载完成
		chromedp.WaitReady("body", chromedp.ByQuery),
		// 注入打印样式
		chromedp.ActionFunc(func(ctx context.Context) error {
			script := fmt.Sprintf(`
                const style = document.createElement('style');
                style.type = 'text/css';
                style.media = 'print';
                style.innerHTML = %q;
                document.head.appendChild(style);
            `, printStyles)

			_, err := page.AddScriptToEvaluateOnNewDocument(script).Do(ctx)
			return err
		}),
		// 等待样式应用
		chromedp.Sleep(2*time.Second),
		// 生成PDF
		chromedp.ActionFunc(func(ctx context.Context) error {
			var err error
			pdfBuffer, _, err = page.PrintToPDF().
				WithPrintBackground(opts.PrintBackground).
				WithPreferCSSPageSize(true).
				WithScale(opts.Scale).
				WithPaperWidth(opts.PaperWidth).
				WithPaperHeight(opts.PaperHeight).
				WithMarginTop(opts.MarginTop).
				WithMarginBottom(opts.MarginBottom).
				WithMarginLeft(opts.MarginLeft).
				WithMarginRight(opts.MarginRight).
				WithLandscape(opts.Landscape).
				WithDisplayHeaderFooter(opts.DisplayHeaderFooter).
				Do(ctx)
			return err
		}),
		chromedp.Evaluate(`document.charset = 'UTF-8';`, nil),
	); err != nil {
		return fmt.Errorf("failed to generate PDF: %v", err)
	}

	// 保存PDF文件
	if err := os.WriteFile(outputPath, pdfBuffer, 0644); err != nil {
		return fmt.Errorf("failed to save PDF: %v", err)
	}

	return nil
}

// handlePDFGeneration 处理PDF生成请求
func handlePDFGeneration(w http.ResponseWriter, r *http.Request) {
	// 检查HTTP方法
	if r.Method != http.MethodPost {
		sendJSONResponse(w, http.StatusMethodNotAllowed, PDFResponse{
			Status: "error",
			Error:  "Method not allowed",
		})
		return
	}

	// 解析请求体
	var req PDFRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		sendJSONResponse(w, http.StatusBadRequest, PDFResponse{
			Status: "error",
			Error:  "Invalid request body",
		})
		return
	}

	// 验证URL
	if req.URL == "" {
		sendJSONResponse(w, http.StatusBadRequest, PDFResponse{
			Status: "error",
			Error:  "URL is required",
		})
		return
	}

	// 使用默认选项（如果未提供）
	if req.Options == nil {
		req.Options = defaultPDFOpts
	}

	// 生成默认文件名（如果未提供）
	if req.Filename == "" {
		req.Filename = fmt.Sprintf("report_%s.pdf", time.Now().Format("20060102150405"))
	}

	// 确保文件名以.pdf结尾
	if !strings.HasSuffix(req.Filename, ".pdf") {
		req.Filename += ".pdf"
	}

	// 创建新任务
	task := taskManager.CreateTask(req)

	// 异步处理PDF生成
	go func() {
		outputPath := filepath.Join("output", req.Filename)
		taskManager.UpdateTaskStatus(task.ID, TaskStatusRunning, nil)

		if err := generatePDF(req.URL, outputPath, req.Options); err != nil {
			taskManager.UpdateTaskStatus(task.ID, TaskStatusFailed, err)
			return
		}

		taskManager.UpdateTaskStatus(task.ID, TaskStatusCompleted, nil)
	}()

	// 返回任务ID
	sendJSONResponse(w, http.StatusAccepted, map[string]interface{}{
		"status":  "accepted",
		"message": "PDF generation started",
		"taskId":  task.ID,
	})
}

// handlePDFDownload 处理PDF下载请求
func handlePDFDownload(w http.ResponseWriter, r *http.Request) {
	// 从URL路径中获取文件名
	filename := filepath.Base(r.URL.Path)
	filePath := filepath.Join("output", filename)

	// 检查文件是否存在
	if _, err := os.Stat(filePath); os.IsNotExist(err) {
		http.NotFound(w, r)
		return
	}

	// 设置响应头
	w.Header().Set("Content-Type", "application/pdf")
	w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%s", filename))
	http.ServeFile(w, r, filePath)
}

// handleTaskStatus 处理任务状态查询请求
func handleTaskStatus(w http.ResponseWriter, r *http.Request) {
	// 从URL路径中获取任务ID
	taskID := strings.TrimPrefix(r.URL.Path, "/task-status/")

	// 获取任务状态
	task, exists := taskManager.GetTask(taskID)
	if !exists {
		sendJSONResponse(w, http.StatusNotFound, PDFResponse{
			Status: "error",
			Error:  "Task not found",
		})
		return
	}

	sendJSONResponse(w, http.StatusOK, task)
}

// sendJSONResponse 发送JSON格式的响应
func sendJSONResponse(w http.ResponseWriter, status int, response interface{}) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	json.NewEncoder(w).Encode(response)
}

// CreateTask 创建新的PDF生成任务
func (tm *TaskManager) CreateTask(req PDFRequest) *Task {
	tm.tasksMux.Lock()
	defer tm.tasksMux.Unlock()

	// 创建新任务
	task := &Task{
		ID:        req.Filename,
		Request:   req,
		Status:    TaskStatusPending,
		CreatedAt: time.Now(),
		UpdatedAt: time.Now(),
	}

	// 保存任务
	tm.tasks[task.ID] = task
	return task
}

// GetTask 获取指定ID的任务
func (tm *TaskManager) GetTask(taskID string) (*Task, bool) {
	tm.tasksMux.RLock()
	defer tm.tasksMux.RUnlock()

	task, exists := tm.tasks[taskID]
	return task, exists
}

// UpdateTaskStatus 更新任务状态
func (tm *TaskManager) UpdateTaskStatus(taskID string, status TaskStatus, err error) {
	tm.tasksMux.Lock()
	defer tm.tasksMux.Unlock()

	// 更新任务状态
	if task, exists := tm.tasks[taskID]; exists {
		task.Status = status
		task.UpdatedAt = time.Now()
		if err != nil {
			task.Error = err.Error()
		}
	}
}
